from spatialdata import SpatialData

# Load tables lazily (keeps large tables out of memory)
sdata = SpatialData.read("large_dataset.zarr", lazy=True)

# Access table - data is loaded on-demand
table = sdata.tables["my_table"]
# table.X is now backed by Dask/Zarr, not loaded into memory

import numpy as np
from scipy import sparse
import anndata as ad
import psutil
import tempfile
from pathlib import Path

# Create synthetic sparse data (100k pixels x 100k m/z bins, 3000 peaks/pixel)
rng = np.random.default_rng(42)
n_pixels, n_mz, peaks_per_pixel = 100000, 100000, 3000
nnz = int(n_pixels * peaks_per_pixel)

X = sparse.csc_matrix(
    (rng.lognormal(7, 1.5, nnz).astype(np.float32),
     (rng.integers(0, n_pixels, nnz), rng.integers(0, n_mz, nnz))),
    shape=(n_pixels, n_mz)
)
print(f"Sparse matrix: {X.shape}, nnz={X.nnz:,}")

# Create and write AnnData
grid_width = int(np.ceil(np.sqrt(n_pixels)))
adata = ad.AnnData(
    X=X,
    obs={"x": np.arange(n_pixels) % grid_width, "y": np.arange(n_pixels) // grid_width},
)
adata.obs_names = [f"pixel_{i}" for i in range(n_pixels)]
adata.var_names = [f"mz_{i}" for i in range(n_mz)]

zarr_path = Path(tempfile.mkdtemp()) / "test.zarr"
adata.write_zarr(str(zarr_path))

# Compare lazy vs eager loading
from anndata.experimental import read_lazy
from anndata import read_zarr

def get_mem():
    return psutil.Process().memory_info().rss / 1e6

# Lazy loading
mem_before = get_mem()
adata_lazy = read_lazy(str(zarr_path))
mem_lazy = get_mem() - mem_before

# Eager loading  
mem_before = get_mem()
adata_eager = read_zarr(str(zarr_path))
mem_eager = get_mem() - mem_before

print(f"Lazy:  +{mem_lazy:.1f} MB, X type: {type(adata_lazy.X).__name__}")
print(f"Eager: +{mem_eager:.1f} MB, X type: {type(adata_eager.X).__name__}")
print(f"Memory savings: {(1 - mem_lazy/mem_eager)*100:.0f}%")